{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "F:\\Ana\\anaconda3\\envs\\neo4j\\Lib\\site-packages\\torch\\utils\\_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import transforms, datasets, models\n",
    "from sklearn.metrics import accuracy_score, precision_score, f1_score, classification_report\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from torch.utils.data import random_split\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据集路径\n",
    "# 数据集路径\n",
    "data_dir = r'F:\\Ana\\python\\projectNeo4j\\Finaly_AlexNet\\alex\\pictures_set\\Pictures_data\\photos'  # 替换为你的数据集路径\n",
    "classes = ['butterfly', 'car', 'cat', 'dog',  'horse', 'landscapes']  # 替换为你的类别名称\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 图像尺寸\n",
    "img_width, img_height = 224, 224\n",
    "batch_size = 32\n",
    "num_epochs = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据增强和预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((img_width, img_height)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet 标准化\n",
    "])\n",
    "\n",
    "# 加载数据集\n",
    "dataset = datasets.ImageFolder(root=data_dir, transform=transform)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 划分训练集和验证集\n",
    "train_size = int(0.8 * len(dataset))\n",
    "val_size = len(dataset) - train_size\n",
    "train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/resnet18-f37072fd.pth\" to else\\hub\\checkpoints\\resnet18-f37072fd.pth\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 44.7M/44.7M [02:58<00:00, 263kB/s]\n",
      "F:\\Ana\\anaconda3\\envs\\neo4j\\Lib\\site-packages\\torch\\optim\\lr_scheduler.py:62: UserWarning: The verbose parameter is deprecated. Please use get_last_lr() to access the learning rate.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['TORCH_HOME'] = 'else'\n",
    "\n",
    "# 创建 DataLoader\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "import torchvision.models as models\n",
    "\n",
    "# 使用最新的权重参数加载预训练的 ResNet18 模型\n",
    "model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)\n",
    "num_ftrs = model.fc.in_features\n",
    "model.fc = nn.Linear(num_ftrs, len(classes))  # 修改最后一层以适应类别数量\n",
    "\n",
    "# 将模型移动到 GPU（如果可用）\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = model.to(device)\n",
    "\n",
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3, verbose=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 早停回调\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=5, verbose=False, delta=0):\n",
    "        self.patience = patience\n",
    "        self.verbose = verbose\n",
    "        self.delta = delta\n",
    "        self.best_score = None\n",
    "        self.early_stop = False\n",
    "        self.counter = 0\n",
    "\n",
    "    def __call__(self, val_loss):\n",
    "        if self.best_score is None:\n",
    "            self.best_score = val_loss\n",
    "        elif val_loss > self.best_score + self.delta:\n",
    "            self.counter += 1\n",
    "            if self.verbose:\n",
    "                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_score = val_loss\n",
    "            self.counter = 0\n",
    "\n",
    "early_stopping = EarlyStopping(patience=5, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50, Training Loss: 0.5280, Validation Loss: 1.0460, Accuracy: 0.7458, Precision: 0.8054, F1 Score: 0.7210\n",
      "Epoch 2/50, Training Loss: 0.3558, Validation Loss: 0.4084, Accuracy: 0.8658, Precision: 0.8797, F1 Score: 0.8649\n",
      "Epoch 3/50, Training Loss: 0.2906, Validation Loss: 0.4275, Accuracy: 0.8591, Precision: 0.8771, F1 Score: 0.8577\n",
      "EarlyStopping counter: 1 out of 5\n",
      "Epoch 4/50, Training Loss: 0.2472, Validation Loss: 0.2849, Accuracy: 0.9061, Precision: 0.9073, F1 Score: 0.9062\n",
      "Epoch 5/50, Training Loss: 0.1923, Validation Loss: 0.5214, Accuracy: 0.8368, Precision: 0.8539, F1 Score: 0.8351\n",
      "EarlyStopping counter: 1 out of 5\n",
      "Epoch 6/50, Training Loss: 0.1636, Validation Loss: 0.4043, Accuracy: 0.8701, Precision: 0.8851, F1 Score: 0.8685\n",
      "EarlyStopping counter: 2 out of 5\n",
      "Epoch 7/50, Training Loss: 0.1381, Validation Loss: 0.3789, Accuracy: 0.8865, Precision: 0.8927, F1 Score: 0.8868\n",
      "EarlyStopping counter: 3 out of 5\n",
      "Epoch 8/50, Training Loss: 0.1245, Validation Loss: 0.4046, Accuracy: 0.8832, Precision: 0.8858, F1 Score: 0.8838\n",
      "EarlyStopping counter: 4 out of 5\n",
      "Epoch 9/50, Training Loss: 0.0489, Validation Loss: 0.2625, Accuracy: 0.9273, Precision: 0.9274, F1 Score: 0.9273\n",
      "Epoch 10/50, Training Loss: 0.0215, Validation Loss: 0.2805, Accuracy: 0.9264, Precision: 0.9271, F1 Score: 0.9265\n",
      "EarlyStopping counter: 1 out of 5\n",
      "Epoch 11/50, Training Loss: 0.0160, Validation Loss: 0.2699, Accuracy: 0.9284, Precision: 0.9292, F1 Score: 0.9286\n",
      "EarlyStopping counter: 2 out of 5\n",
      "Epoch 12/50, Training Loss: 0.0110, Validation Loss: 0.2865, Accuracy: 0.9329, Precision: 0.9335, F1 Score: 0.9330\n",
      "EarlyStopping counter: 3 out of 5\n",
      "Epoch 13/50, Training Loss: 0.0093, Validation Loss: 0.2967, Accuracy: 0.9316, Precision: 0.9318, F1 Score: 0.9315\n",
      "EarlyStopping counter: 4 out of 5\n",
      "Epoch 14/50, Training Loss: 0.0062, Validation Loss: 0.2954, Accuracy: 0.9304, Precision: 0.9314, F1 Score: 0.9306\n",
      "EarlyStopping counter: 5 out of 5\n",
      "Early stopping triggered\n"
     ]
    }
   ],
   "source": [
    "# 训练模型\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs):\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        running_loss = 0.0\n",
    "        for inputs, labels in train_loader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            running_loss += loss.item()\n",
    "\n",
    "        # 验证阶段\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        y_true, y_pred = [], []\n",
    "        with torch.no_grad():\n",
    "            for inputs, labels in val_loader:\n",
    "                inputs, labels = inputs.to(device), labels.to(device)\n",
    "                outputs = model(inputs)\n",
    "                loss = criterion(outputs, labels)\n",
    "                val_loss += loss.item()\n",
    "                _, predicted = torch.max(outputs, 1)\n",
    "                y_true.extend(labels.cpu().numpy())\n",
    "                y_pred.extend(predicted.cpu().numpy())\n",
    "\n",
    "        # 计算验证集损失和指标\n",
    "        val_loss /= len(val_loader)\n",
    "        accuracy = accuracy_score(y_true, y_pred)\n",
    "        precision = precision_score(y_true, y_pred, average='weighted')\n",
    "        f1 = f1_score(y_true, y_pred, average='weighted')\n",
    "\n",
    "        print(f'Epoch {epoch + 1}/{num_epochs}, '\n",
    "              f'Training Loss: {running_loss / len(train_loader):.4f}, '\n",
    "              f'Validation Loss: {val_loss:.4f}, '\n",
    "              f'Accuracy: {accuracy:.4f}, '\n",
    "              f'Precision: {precision:.4f}, '\n",
    "              f'F1 Score: {f1:.4f}')\n",
    "\n",
    "        # 早停检查\n",
    "        early_stopping(val_loss)\n",
    "        if early_stopping.early_stop:\n",
    "            print(\"Early stopping triggered\")\n",
    "            break\n",
    "\n",
    "        # 调整学习率\n",
    "        scheduler.step(val_loss)\n",
    "\n",
    "# 开始训练\n",
    "train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs)\n",
    "\n",
    "# 保存模型\n",
    "torch.save(model.state_dict(), 'image_classification_model_2.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "   butterfly       0.96      0.94      0.95       723\n",
      "         car       0.92      0.92      0.92       717\n",
      "         cat       0.93      0.92      0.93       746\n",
      "         dog       0.94      0.89      0.91       786\n",
      "       horse       0.98      0.98      0.98       667\n",
      "  landscapes       0.87      0.93      0.90       803\n",
      "\n",
      "    accuracy                           0.93      4442\n",
      "   macro avg       0.93      0.93      0.93      4442\n",
      "weighted avg       0.93      0.93      0.93      4442\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 加载模型并评估\n",
    "model.load_state_dict(torch.load('image_classification_model_2.pth'))\n",
    "model.eval()\n",
    "y_true, y_pred = [], []\n",
    "with torch.no_grad():\n",
    "    for inputs, labels in val_loader:\n",
    "        inputs, labels = inputs.to(device), labels.to(device)\n",
    "        outputs = model(inputs)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        y_true.extend(labels.cpu().numpy())\n",
    "        y_pred.extend(predicted.cpu().numpy())\n",
    "\n",
    "# 打印分类报告\n",
    "print(classification_report(y_true, y_pred, target_names=classes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
